package cc.spider.web.queue;

import cc.spider.core.queue.AbsUniqueQueue;
import cc.spider.core.queue.ITaskQueue;
import cc.spider.model.TaskInfo;
import cc.spider.web.model.http.Request;
import org.apache.commons.lang3.StringUtils;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;

import java.io.*;

/**
 * 基于redis实现的优先级去重队列
 */
public class RedisUniqueQueue extends AbsUniqueQueue {

    private JedisPool jedisPool = new JedisPool();

    private static final String QUEUE_PREFIX = "enqueue_";

    private static final String DEQUEUE_PREFIX = "dequeue_";

    private SerializeUtil<TaskInfo> serializeUtil = new SerializeUtil<>();
    private String uuid;

    public RedisUniqueQueue(String uuid) {
        this.uuid = uuid;
    }

    private boolean add(String uuKey) {
        Jedis jedis = jedisPool.getResource();
        boolean res = jedis.sadd(DEQUEUE_PREFIX + uuid, uuKey) > 0;
        jedisPool.returnResource(jedis);
        return res;
    }

    @Override
    public TaskInfo dequeue() {
        Jedis jedis = jedisPool.getResource();
        String content = jedis.rpop(QUEUE_PREFIX + uuid);
        jedisPool.returnResource(jedis);
        return null == content ? null : serializeUtil.unserialize(content);
    }

    @Override
    public boolean isUnique(TaskInfo taskInfo) {
        return (taskInfo instanceof Request && StringUtils.containsIgnoreCase("POST", ((Request) taskInfo).getMethod())) || this.add(taskInfo.getUuKey());
    }

    @Override
    public boolean enqueueWhenUnique(TaskInfo taskInfo) {
        Jedis jedis = jedisPool.getResource();
        boolean res = jedis.rpush(QUEUE_PREFIX + uuid, serializeUtil.serialize(taskInfo)) > 0;
        if(!res){
           jedis.srem(DEQUEUE_PREFIX + uuid, taskInfo.getUuKey());
        }
        jedisPool.returnResource(jedis);
        return res;
    }

    private static class SerializeUtil<T> {
        public String serialize(T object) {
            ObjectOutputStream oos = null;
            ByteArrayOutputStream baos = null;
            try {
                //序列化
                baos = new ByteArrayOutputStream();
                oos = new ObjectOutputStream(baos);
                oos.writeObject(object);
                return baos.toString("ISO-8859-1");
            } catch (Exception e) {
                e.printStackTrace();
                return null;
            }
        }

        public T unserialize(String bytes) {
            ByteArrayInputStream bais = null;
            try {
                //反序列化
                bais = new ByteArrayInputStream(bytes.getBytes("ISO-8859-1"));
                ObjectInputStream ois = new ObjectInputStream(bais);
                return (T) ois.readObject();
            } catch (Exception e) {
                e.printStackTrace();
                return null;
            }
        }
    }

}
